package com.geeguo.ebuilder.core.feign.utils;

import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.CompletableFuture;
import java.util.function.Supplier;

/**
 * 异步 Feign 请求工具。
 * 把当前 Request 的 Header 存入新线程中，在新线程中调用 Feign 请求，FeignConfig 中取当前 headers发送到新服务
 */
@Slf4j
@Component
@EnableAsync
public class FeignHolder {

    private static ThreadLocal<Map<String, String>> feignHeader = ThreadLocal.withInitial(HashMap::new);

    /**
     * 设置 Feign 请求头
     */
    public static void set(Map<String, String> headers) {
        if (headers == null) {
            throw new IllegalArgumentException("Headers cannot be null.");
        }
        // 防御性拷贝
        feignHeader.set(new HashMap<>(headers));
    }

    /**
     * 获取当前线程的 Feign 请求头
     *
     * @return 不可修改的请求头 Map，可能为空但不会为 null。
     */
    public static Map<String, String> get() {
        return Collections.unmodifiableMap(feignHeader.get());
    }

    /**
     * 获取当前线程的 Feign 请求头的副本
     *
     * @return 新的 HashMap 副本，可修改。
     */
    public static Map<String, String> getCopy() {
        return new HashMap<>(feignHeader.get());
    }

    /**
     * 清除当前线程的Feign请求头
     */
    public static void clear() {
        feignHeader.remove();
    }

    /**
     * 异步执行 Feign 请求
     *
     * @param headers 请求头
     * @param task 要执行的任务
     * @return CompletableFuture以便链式调用
     */
    @Async
    public CompletableFuture<Void> asyncFeign(Map<String, String> headers, Runnable task) {
        if (headers == null || task == null) {
            throw new IllegalArgumentException("Headers and task cannot be null.");
        }
        try {
            set(headers);
            task.run();
            return CompletableFuture.completedFuture(null);
        } catch (Exception e) {
            log.error("Async Feign request failed.", e);
            return CompletableFuture.failedFuture(e);
        } finally {
            clear();
        }
    }

    /**
     * 带返回值的异步Feign请求
     *
     * @param headers 请求头
     * @param supplier 提供返回值的任务
     * @param <T> 返回值类型
     * @return CompletableFuture包装的结果
     */
    @Async
    public <T> CompletableFuture<T> asyncFeignWithResult(Map<String, String> headers, Supplier<T> supplier) {
        if (headers == null || supplier == null) {
            return CompletableFuture.failedFuture(new IllegalArgumentException("Headers and supplier cannot be null"));
        }
        try {
            set(headers);
            T result = supplier.get();
            return CompletableFuture.completedFuture(result);
        } catch (Exception e) {
            log.error("Async Feign request with result failed.", e);
            return CompletableFuture.failedFuture(e);
        } finally {
            clear();
        }
    }
}